import csv
import warnings
warnings.filterwarnings('ignore')
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score
from sklearn.naive_bayes import  MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split


def getSentiment_data():
    sentiment_data = []
    with open('./nlp_tiktok_date/train1.csv', 'r', encoding='utf-8') as readerFile:
        reader = csv.reader(readerFile)
        for row in reader:
            sentiment_data.append(row)
    return sentiment_data

def model_train():
    sentiment_data = getSentiment_data()
    df = pd.DataFrame(sentiment_data, columns=['txt', 'label'])

    train_date,test_data=train_test_split(df,test_size=0.2,random_state=43)

    # 原始文本转为tf-idf的特征矩阵
    vectorizer = TfidfVectorizer()
    X_train = vectorizer.fit_transform(train_date['txt'])
    y_train = train_date['label']
    X_test = vectorizer.transform(test_data['txt'])
    y_test = test_data['label']

    classifier = MultinomialNB()
    classifier.fit(X_train, y_train)
    y_pred = classifier.predict(X_test)
    #计算正确率
    accuracy = accuracy_score(y_test, y_pred)
    print(accuracy)
    print(test_data,y_pred)
    import joblib # 新版本 Scikit-learn
    joblib.dump(classifier, "./model/train_model.m")




if __name__ == "__main__":
    model_train()